# $Id: sshlib.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

import base64
import binascii
import getpass
import os
import socket
import warnings

with warnings.catch_warnings():
   warnings.filterwarnings("ignore", category = DeprecationWarning)
   import paramiko

import log
import sshshell
import stdlib

gethex = lambda key: binascii.hexlify(key)

class SSHException(Exception):
   pass

class SSHClient(object):
   def __init__(self, **kwargs):
      self._err = SSHException()
      self._sock = None
      self._host = None
      self._port = 22
      self._ssh = None
      self._channel = None
      self._username = None
      self._password = None
      self._keychangeok = False
      self._log = None
      self._shell = None
      self._sftp = None
      self._auth = False

      if "host" in kwargs:
         self._host = kwargs['host'].strip()
      else:
         self._error("No hostname provided!")

      if "user" in kwargs:
         self._username = kwargs['user'].strip()
      else:
         self._username = getpass.getuser()

      if "port" in kwargs:
         try:
            self._port = int(kwargs['port'])
         except ValueError, message:
            self._error(message)

      if "passwd" in kwargs:
         if kwargs['passwd']:
            self._password = kwargs['passwd'].strip()

      if "keychangeok" in kwargs:
         self._keychangeok = bool(kwargs['keychangeok'])

      verbose = False
      if "verbose" in kwargs:
         verbose = bool(kwargs["verbose"])

      if "log" in kwargs:
         if kwargs['log'] == 'syslog':
            self._log = log.Log(self._name, str(self._host))

      if self._log is None:
         if verbose:
            self._log = log.StreamLog(self._name)
         else:
            self._log = log.NullLog(self._name)

   def __enter__(self):
      return self

   def __exit__(self, type, value, traceback):
      self.Close()
      if type and value:
         raise type, value

   @property
   def _name(self):
      return stdlib.my_name(self)

   def __str__(self):
      retval = ""
      if self._sock:
         retval = "ssh socket to %s on port %d" %(self._host, self._port)

      return retval.encode('utf-8')

   def __repr__(self):
      return repr(str(self))

   def __unicode__(self):
      return unicode(str(self))

   def __int__(self):
      retval = -1
      try:
         retval = self._sock.fileno()
      except socket.error, (error, message):
         pass

      return int(retval)

   def _error(self, fmt, *args):
      raise SSHException, stdlib.sprintf(fmt, *args)

   def Close(self):
      if self._ssh:
         self._ssh.close()
         self._ssh = None

      if self._sock:
         self._sock.close()
         self._sock = None

      if self._shell:
         self._shell.Close()
         self._shell = None

   def Connect(self, family = None):
      sockList = []
      mesg = None

      addrfam = socket.AF_UNSPEC
      if family == 4:
         addrfam = socket.AF_INET
      elif family == 6:
         addrfam = socket.AF_INET6

      try:
         sockList = socket.getaddrinfo(self._host, self._port, addrfam,
               socket.SOCK_STREAM)
      except socket.gaierror, (error, message):
         self._error("getaddrinfo: error %d, message: %s", error, message)

      for res in sockList:
         af, socktype, proto, canonname, sa = res
         try:
            self._sock = socket.socket(af, socktype, proto)
         except socket.error, (error, message):
            self._sock = None
            continue

         if self._sock:
            self._log("socket fd: %d", self._sock.fileno())

         try:
            self._sock.connect(sa)
         except socket.error, (error, message):
            mesg = message
            self._sock.close()
            self._sock = None
            continue
         else:
            try:
               (host, port) = socket.getnameinfo(sa, 1)
               if host and port:
                  self._log("connected to %s (%s) on %s port",
                        self._host, host, port)
            except:
               self._log("connected to %s on %s port", self._host, port)

         break

      if self._sock is None:
         self._error("cannot connect to %s on tcp port %d: %s", self._host,
               self._port, mesg)

      self._ssh = paramiko.Transport(self._sock)
      try:
         self._ssh.start_client()
      except paramiko.SSHException:
         self._error("SSH negotiation failed!")

      self._log("SSH negotiation complete!")

      knownhosts = os.path.join(os.environ['HOME'], '.ssh', 'known_hosts')
      self._log("Known hosts file: %s", knownhosts)

      keys = {}
      try:
         keys = paramiko.util.load_host_keys(knownhosts)
      except IOError, message:
         keys = {}

      key = self._ssh.get_remote_server_key()
      self._log("%s key: %s", self._host, gethex(key.get_fingerprint()))

      if self._host not in keys:
         self._log("WARNING: Unknown host key for %s", str(self._host))
      elif not keys[self._host].has_key(key.get_name()):
         self._log("WARNING: Unknown host key for %s", str(self._host))
      elif keys[self._host][key.get_name()] != key:
         if self._keychangeok:
            self._log("WARNING: Unknown host key for %s", str(self._host))
         else:
            self._error("WARNING: Unknown host key for %s", str(self._host))

   def pkauth(self):
      retval = False
      try:
         agent = paramiko.Agent()
      except socket.error, (error, message):
         self._log.info("Cannot open your authentication agent!")
      else:
         keys = agent.get_keys()
         for key in keys:
            try:
               self._ssh.auth_publickey(self._username, key)
            except paramiko.SSHException:
               self._log.info("public key agent forwarding failed!")
            except TypeError, message:
               self._log.info(message)
               pass
            else:
               retval = True

      return retval

   def keyauth(self, password = None):
      retval = False
      path = os.path.join(os.environ['HOME'], '.ssh', 'id_rsa')
      self._log("private key %s", path)
      self._log("public key %s.pub", path)
      if os.path.exists(path):
         try:
            if password:
               key = paramiko.RSAKey.from_private_key_file(path, password)
            else:
               key = paramiko.RSAKey.from_private_key_file(path)
         except paramiko.PasswordRequiredException:
            password = getpass.getpass("Enter passphrase for key %s: " %path)
            retval = self.keyauth(password.strip())
         except paramiko.SSHException, message:
            self._log.info(message)
         else:
            try:
               self._ssh.auth_publickey(self._username, key)
            except paramiko.SSHException:
               self._log.info("Cannot authenticate using public key!")
            else:
               retval = True

      return retval

   def passauth(self, error = False):
      retval = False
      try:
         if not self._password:
            self._password = getpass.getpass("Password for %s@%s: " %(
               self._username, self._host)).strip()

         self._ssh.auth_password(self._username, self._password)
      except paramiko.AuthenticationException, message:
         if error:
            self._error("Cannot login to %s on %s: %s", self._username,
                  self._host, message)
         else:
            self._log("Cannot login to %s on %s: %s", self._username,
                  self._host, message)
      else:
         retval = True

      return retval

   def Auth(self, Password = False):
      retval = False
      if self._ssh:
         try:
            if Password:
               retval = self.passauth(True)
            else:
               if self.pkauth():
                  retval = True
               elif self.keyauth():
                  retval = True
               elif self.passauth():
                  retval = True
         except SSHException, message:
            self._error(message)

         if retval:
            self._auth = True

      return retval

   def GetChannel(self, **kwargs):
      if self._ssh:
         exfile = None
         sendstdout = True

         if "exfile" in kwargs:
            if kwargs['exfile'] is not None:
               exfile = kwargs['exfile'].strip()

         if "sendstdout" in kwargs:
            sendstdout = bool(kwargs['sendstdout'])

         if not self._channel:
            self._channel = self._ssh.open_session()
            self._channel.get_pty()
            self._channel.invoke_shell()

         try:
            self._shell = sshshell.Terminal(self._channel)
         except sshshell.TerminalError, message:
            self._error(message)

         if exfile:
            try:
               self._shell.SendFile(exfile, sendstdout)
            except sshshell.TerminalError, message:
               self._error(message)
         else:
            self._shell.GetShell()

   def sftp(self):
      if self._auth:
         self._sftp = paramiko.SFTPClient.from_transport(self._ssh)

   def Get(self, Remote, local = None):
      if not local:
         local = Remote

      self._sftp.get(Remote, local)

   def Put(self, Local, remote = None):
      if not remote:
         remote = Local

      self._sftp.put(remote, Local)

